# -*- coding: utf-8 -*-
"""
Created on Wed May 12 20:45:07 2021

@author: 31943
"""
from facenet_pytorch import MTCNN, InceptionResnetV1
from torchvision import transforms
import torch
import os
import time
from PIL import Image
import shutil
import socket
import select
from threading import Thread


# 全局变量定义与超参数
# MT-CNN 网络负责检测人脸 ， ResNet 负责对检测出的人脸进行 512 维特征的提取
workers = 0 if os.name == 'nt' else 4
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
MtCnn = MTCNN(keep_all=True, device=device)
ResNet = InceptionResnetV1(pretrained='vggface2').eval().to(device)
transform = transforms.Compose([
        # 按照比例把图像最小的一个边长放缩到256，另一边按照相同比例放缩。
        transforms.Resize(256),
    ])
# 线程池
threads = []


# 图片缓存的固定路径
fixed_path = "D:/JAVA-competition/test/wisdom/"
aim_path = ""
cache_path = ""


# 转码utf-8
def utf8(text):
    return str(text, 'utf-8')


# 转码gbk
def gbk(text):
    return text.decode("gbk", "ignore")


# 如果有新的人物加入, 保存 names , embeddings 文件
def save(embed, name, path):
    torch.save(embed, path + "embeddings.pt")  # 当然也可以保存在一个文件
    torch.save(name, path + "names.pt")


# 加载模型---如果有新人加入时需要重加载
def loadModel(path):
    name = torch.load(path + "names.pt")
    embed = torch.load(path + "embeddings.pt").to(device).requires_grad_(False)
    return name, embed


# 如果用户是第一次使用, 目录下无模型，应首先新建模型，以便于后续调用
def firstTimeLoadModel(picture, path, person):
    aligned = []
    names = []
    embeddings = []
    # 获取人脸区域
    faces, boxes = MtCnn(picture)
    if faces is not None:
        aligned.append(faces)
        names.append('person0')
        aligned = torch.stack(aligned).to(device)
        # 提取人脸特征向量，每个向量的长度为512
        for p, box in enumerate(boxes):
            aligned = ResNet(faces[p].unsqueeze(0).to(device))  # 提取人脸特征向量
        aligned = aligned.detach().numpy()
        embeddings.append(aligned[0] * 0)
        embeddings = torch.Tensor(embeddings)
        embeddings.requires_grad_(False)
        # 生成person0文件夹
        # 理论上这里应该用线程锁的，但为了方便，暂时使用判断来解决
        if not os.path.exists(path + "person0"):
            os.makedirs(path + "person0")
        # 把这第一张单人照的图片移动到对应目录下
        # shutil.move(cache_path + person, aim_path + 'person0/')
        # 保存模型
        torch.save(embeddings, path + "embeddings.pt")
        torch.save(names, path + "names.pt")


# 对输入的新照片进行人脸检测和识别
def detect_photo(picture, person):
    # 这里的模型加载需要考虑一下按需加载，如果人的分类多起来了会相当耗时
    faces, boxes = MtCnn(picture)
    # 由于是已经分类好的people类，因此基本必有人，若识别不出人脸
    # 说明人脸像素过低，一般此种情况是由群体照导致的(极远照则被认定为风景)
    if boxes is None or len(boxes) >= 4:
        return "more", "old"
    if len(boxes) == 1:
        if not os.path.exists(aim_path + "names.pt") or not os.path.exists(aim_path + "embeddings.pt"):
            firstTimeLoadModel(picture, aim_path, person)
            return "OK", "new"
    if len(boxes) == 2:
        return "two", "old"
    if len(boxes) == 3:
        return "three", "old"
    for p, box in enumerate(boxes):
        names, embeddings = loadModel(aim_path)
        face_embedding = ResNet(faces[p].unsqueeze(0).to(device))  # 提取人脸特征向量
        # 计算与库中 保存的所有的 人脸特征数据的欧式距离(相似度)
        similarity = [(face_embedding - embeddings[embed]).norm().item() for embed in range(embeddings.size()[0])]
        threshold = 1.0001 - 0.1*(len(similarity) - 1)/(len(similarity))
        if min(similarity) < threshold:    # 欧氏距离阈值
            index = similarity.index(min(similarity))  # 对应的索引就是判断的人脸
            person_name = names[index]  # 对应的人脸
            # 误删文件夹时
            if not os.path.exists(aim_path + person_name):
                os.makedirs(aim_path + person_name)
            # 更新一下 embeddings 文件， 权重按照 0.6:0.4
            embeddings[index] = embeddings[index] * 0.6 + face_embedding[0] * 0.4
            save(embeddings, names, aim_path)
            return person_name, "old"
        else:
            # 与库中的所有的人脸特征都不相似，新建人脸文件夹
            if not os.path.exists(aim_path + 'person' + str(len(names))):
                os.makedirs(aim_path + 'person' + str(len(names)))
            # 与库中的所有的人脸特征都不相似，新建人脸类别
            names.append('person' + str(len(names)))
            # 更新一下 embeddings 文件
            embeddings = torch.cat([embeddings, face_embedding], dim=0)  # 在dim=0 处拼接
            save(embeddings, names, aim_path)
            return 'person' + str(len(names) - 1), "new"


# 图片读入与处理函数
def imgReadAndSolve(*people):
    for person in people:
        img = Image.open(cache_path + person)
        img = transform(img)
        result, status = detect_photo(img, person)
        if result == "two":
            # 是2人照, 将此图片移动至two文件夹, 填充报文
            if result in old:
                old[result].append(person)
            else:
                old[result] = [person]
            # shutil.move(cache_path + person, aim_path + 'two/')
        elif result == "three":
            # 是3人照, 将此图片移动至three文件夹, 填充报文
            if result in old:
                old[result].append(person)
            else:
                old[result] = [person]
            # shutil.move(cache_path + person, aim_path + 'three/')
        elif result == "more":
            # 是4plus人照, 将此图片移动至more文件夹, 填充报文
            if result in old:
                old[result].append(person)
            else:
                old[result] = [person]
            # shutil.move(cache_path + person, aim_path + 'more/')
        elif result == "OK":
            # 是第一次生成模型+移动的返回结果
            new['person0'] = [person]
        elif result.__len__() > 6:
            # 是个人照
            if result in new:
                new[result].append(person)
            else:
                if status == "new":
                    new[result] = [person]
                else:
                    if result in old:
                        old[result].append(person)
                    else:
                        old[result] = [person]
            # shutil.move(cache_path + person, aim_path + result + '/')
    # 报文处理


if __name__ == '__main__':
    # select-socket通信建立
    sock = socket.socket()
    sock.bind(('127.0.0.1', 4444))
    sock.setblocking(False)
    sock.listen(10)
    inputs = [sock, ]

    while True:
        r_list, w_list, e_list = select.select(inputs, [], [], 1)
        for event in r_list:
            if event == sock:
                print("新的客户端连接")
                new_sock, addr = event.accept()
                inputs.append(new_sock)
            else:
                data = event.recv(1024)
                if data:
                    # 收到用户的路径名，进行拼接后进行人脸分类
                    # 临时文件夹路径
                    cache_path = fixed_path + utf8(data) + '/people/cache/'
                    # 目标文件夹路径
                    aim_path = fixed_path + utf8(data) + '/people/'
                    if os.path.exists(cache_path):
                        # 如果路径存在
                        files = os.listdir(cache_path)
                        # 若有模型则处理第一张，若无则生成模型
                        if files.__len__() != 0:
                            category = {}  # 报文字典
                            new = {}  # 新增字典
                            old = {}  # 旧字典
                            thread = Thread(target=imgReadAndSolve, args=files[0:files.__len__()])
                            threads.append(thread)
                            thread.setDaemon(True)
                            threads[0].start()
                            threads[0].join()
                            threads = []
                            for key in new:
                                for picture in new[key]:
                                    # 将图片移动至对应文件夹
                                    shutil.move(cache_path + picture, aim_path + key + '/')
                            for key in old:
                                for picture in old[key]:
                                    # 将图片移动至对应文件夹
                                    shutil.move(cache_path + picture, aim_path + key + '/')
                            category = {"new": new, "old": old}
                            print(category)
                            event.send(str(category).encode())
                            event.send("over".encode())
                        else:
                            event.send("over".encode())
                    else:
                        print(cache_path + "路径不存在")
                else:
                    print("客户端断开连接")
                    inputs.remove(event)


